
#include "replication_pad2d_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    ReplicationPad2dTilingData tiling;
    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    auto coreNum = ascendcPlatform.GetCoreNum();

    // 划分成单行tensor，同时附带行数

    uint32_t columns = context->GetInputShape(0)->GetOriginShape().GetDim(context->GetInputShape(0)->GetOriginShape().GetDimNum() -1);
    uint32_t rows = context->GetInputShape(0)->GetOriginShape().GetDim(context->GetInputShape(0)->GetOriginShape().GetDimNum() - 2);
    uint32_t unitCount = 1;
    for(int i = 0; i < context->GetInputShape(0)->GetOriginShape().GetDimNum() - 2; i++) {
        unitCount *= context->GetInputShape(0)->GetOriginShape().GetDim(i);
    }


    // Based on the input length and the number of inputs, the number of bytes of the input data type is obtained
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t typeLength = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);
    uint32_t inputLength = inputNum * typeLength;
    uint32_t inputBytes = inputLength / inputNum;

    // There are a total of 3 shared UB spaces in the input and output. If it's int8, there are 2 more TBUFs
    uint32_t ubDataNumber = 5;
    // The number of 32B data blocks that can be used for each data. DOUBLE BUFFER is already counted here
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
    // uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;
    uint32_t tileDataNum = ((tileBlockNum * BLOCK_SIZE) / inputBytes) / columns * columns;

    // Input data for 32B alignment
    uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
    // There is at least 32B of data on each core, satisfying several settings for several cores. The maximum number of audits is the actual number of audits
    // coreNum = (coreNum <  inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;
    // coreNum = (coreNum >= 1) ? coreNum : 1;
    coreNum = 1;
    // 计算每个核需要处理多少个32B（对于单核应用也即计算所有任务可以分成多少个32B）
    // realInputLength表示一共需要处理多少B的输入数据（未对齐）
    uint32_t realInputLength = rows * columns * typeLength;
    uint32_t everyCoreInputBlockNum = realInputLength / BLOCK_SIZE / coreNum;
    uint32_t tailBlockNum = (realInputLength / BLOCK_SIZE) % coreNum;
    
    // Small chunks are calculated and sliced several times using the number of data on each core
    // 计算小核总共处理多少个数据元素，对于单核场景计算的结果应该是所有的数据元素数
    uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    // tileBlockNum表示考虑到ub之后每个tile最多可以处理多少个数据，smallTileNum就计算出了每个core需要几个tile
    // uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
    // uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
    // Tail block calculation for small chunks of data
    // uint32_t smallTailDataNum = smallCoreDataNum - (tileDataNum * smallTileNum);
    // smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;
    

    // 总共需要处理的数据
    uint32_t realSmallCoreDataNum = unitCount * rows * columns;
    uint32_t smallTailDataNum = realSmallCoreDataNum % (tileDataNum);
    smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;
    uint32_t finalSmallTileNum = realSmallCoreDataNum % tileDataNum == 0 ? realSmallCoreDataNum / tileDataNum : (realSmallCoreDataNum / tileDataNum + 1); 
    
    // The total length of a large block of data is 32B larger than that of a small block of data
    everyCoreInputBlockNum += 1;
    uint32_t bigCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    uint32_t bigTileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalBigTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? bigTileNum : bigTileNum + 1;
    uint32_t bigTailDataNum = bigCoreDataNum - tileDataNum * bigTileNum;
    bigTailDataNum = bigTailDataNum == 0 ? tileDataNum : bigTailDataNum;
    
    tiling.set_smallCoreDataNum(realSmallCoreDataNum);
    tiling.set_bigCoreDataNum(bigCoreDataNum);
    tiling.set_tileDataNum(tileDataNum);
    tiling.set_smallTailDataNum(smallTailDataNum);
    tiling.set_bigTailDataNum(bigTailDataNum);
    tiling.set_finalSmallTileNum(finalSmallTileNum);
    tiling.set_finalBigTileNum(finalBigTileNum);
    tiling.set_tailBlockNum(tailBlockNum);
    tiling.set_columns(columns);
    tiling.set_rows(rows);
    tiling.set_unitCount(unitCount);
    

    printf("unitCount: %u\n", unitCount);
    printf("rows: %u\n", rows);
    printf("columns: %u\n", columns);
    printf("inputNum: %u\n", inputNum);
    printf("finalSmallTileNum: %u\n", finalSmallTileNum);
    printf("smallTailDataNum: %u\n", smallTailDataNum);
    printf("realSmallCoreDataNum: %u\n", realSmallCoreDataNum);
    printf("tileDataNum: %u\n", tileDataNum);
    printf("everyCoreInputBlockNum: %u\n", everyCoreInputBlockNum);
    printf("realInputLength: %u\n", realInputLength);
    
    context->SetBlockDim(coreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ReplicationPad2d : public OpDef {
public:
    explicit ReplicationPad2d(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("paddings")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910");
        this->AICore().AddConfig("ascend310b");
    }
};

OP_ADD(ReplicationPad2d);
}
